import os
import cv2
import numpy as np
import shutil
from insightface.app import FaceAnalysis
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import StandardScaler

class FaceClassifier:
    def __init__(self, source_path, output_path, eps=0.5, min_samples=3):
        """
        初始化人脸分类器
        
        参数:
        source_path: 源图片目录路径
        output_path: 分类结果输出目录
        eps: DBSCAN邻域半径参数（默认0.5）
        min_samples: DBSCAN最小样本数参数（默认3）
        """
        self.source_path = source_path
        self.output_path = output_path
        self.eps = eps
        self.min_samples = min_samples
        
        # 初始化人脸分析模型
        self.app = FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))
        
        # 存储所有人脸特征和对应图片信息
        self.all_face_embeddings = []
        self.all_face_images = []
        self.all_face_filenames = []
        
        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)

    def process_images(self):
        """处理源目录中的所有图片"""
        # 获取所有支持的图片文件
        image_exts = ('.jpg', '.jpeg', '.png', '.bmp', '.tif')
        image_files = [f for f in os.listdir(self.source_path) 
                      if f.lower().endswith(image_exts)]
        
        print(f"发现 {len(image_files)} 张图片需要处理")
        
        # 处理每张图片并收集所有人脸数据
        for i, img_file in enumerate(image_files):
            img_path = os.path.join(self.source_path, img_file)
            print(f"\n处理图片 {i+1}/{len(image_files)}: {img_file}")
            
            # 读取并处理图片
            img = cv2.imread(img_path)
            if img is None:
                print(f"  警告：无法读取图片 {img_file}，跳过")
                continue
            
            # 检测人脸
            faces = self.app.get(img)
            
            if not faces:
                print(f"  未检测到人脸，跳过")
                self._save_unclassified(img, img_file)
                continue
                
            # 处理每张检测到的人脸
            for j, face in enumerate(faces):
                face_embedding = face.embedding
                bbox = face.bbox.astype(int)
                
                # 裁剪人脸区域
                face_img = img[bbox[1]:bbox[3], bbox[0]:bbox[2]]
                
                # 保存人脸数据和相关信息
                self.all_face_embeddings.append(face_embedding)
                self.all_face_images.append(face_img)
                self.all_face_filenames.append(img_file)
        
        # 所有人脸处理完成后进行DBSCAN聚类
        self._cluster_faces()
        
        print("\n处理完成!")
        print(f"共识别出 {len(set(self.cluster_labels)) - (1 if -1 in self.cluster_labels else 0)} 个不同人物")
        print(f"分类结果保存在: {self.output_path}")

    def _cluster_faces(self):
        """使用DBSCAN对所有检测到的人脸进行聚类"""
        if not self.all_face_embeddings:
            print("警告：没有检测到任何人脸")
            return
            
        # 将特征向量转换为numpy数组并归一化
        X = np.array(self.all_face_embeddings)
        X = StandardScaler().fit_transform(X)  # 标准化处理[8](@ref)
        
        # 使用DBSCAN进行聚类
        # 使用余弦距离作为相似度度量[4](@ref)
        dbscan = DBSCAN(eps=self.eps, min_samples=self.min_samples, metric='cosine')
        self.cluster_labels = dbscan.fit_predict(X)
        
        # 为每个聚类创建目录并保存人脸图片
        for i, (label, face_img, filename) in enumerate(zip(
            self.cluster_labels, self.all_face_images, self.all_face_filenames)):
            
            # 噪声点（未聚类成功的点）放入unclassified目录
            if label == -1:
                person_id = "unclassified"
            else:
                person_id = f"Person_{label + 1}"  # 聚类标签从0开始，我们加1使其从1开始
                
            self._save_person_face(person_id, face_img, filename, i)

    def _save_person_face(self, person_id, face_img, original_filename, index):
        """保存人脸图片到对应人物目录"""
        person_dir = os.path.join(self.output_path, person_id)
        os.makedirs(person_dir, exist_ok=True)
        
        # 生成唯一文件名
        base_name = os.path.splitext(original_filename)[0]
        face_filename = f"{base_name}_{index}.jpg"
        face_path = os.path.join(person_dir, face_filename)
        
        # 保存人脸图片
        cv2.imwrite(face_path, face_img)

    def _save_unclassified(self, img, original_filename):
        """保存未分类的图片"""
        unclassified_dir = os.path.join(self.output_path, "unclassified")
        os.makedirs(unclassified_dir, exist_ok=True)
        cv2.imwrite(os.path.join(unclassified_dir, original_filename), img)

if __name__ == "__main__":
    # 配置路径参数
    SOURCE_DIR = "input_images"  # 替换为你的图片目录路径
    OUTPUT_DIR = "output_dbscan"       # 替换为输出目录路径
    
    # 创建分类器并处理图片
    classifier = FaceClassifier(
        source_path=SOURCE_DIR, 
        output_path=OUTPUT_DIR,
        eps=0.7,       # DBSCAN邻域半径参数，可根据数据调整
        min_samples=2  # DBSCAN最小样本数参数，可根据数据调整
    )
    classifier.process_images()
